import numpy as np
import matplotlib
matplotlib.use(backend="TkAgg")
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
import matplotlib.patches as patches

'''
如果信号中确实有 k=3 的成分，这些旋转向量会越累加越长 → 共振 → 幅值大 → 频谱峰值
如果 k 不匹配（例如 k=5），向量方向乱转，累加后会互相抵消 → 接近 0 → 无峰值
'''


# ---------------- Parameters ----------------
N = 64                    # 采样点数

n = np.arange(N)
k1, k2 = 3, 5             # 两个检验频率
interval_ms = 1000         # 每帧间隔 (ms)

# 原始信号：x[n] = cos(2π * 3 * n / N)
x = np.cos(2 * np.pi * 3 * n / N)

# ---------------- Precompute ----------------
# 旋转因子表：W[k, n] = e^{-j*2π*k*n/N}
# 1️⃣ 我们需要计算的频率索引 k 是 0 到 max(k1, k2)
k_values = np.arange(max(k1, k2) + 1)   # 变成 0,1,2,...5

# 2️⃣ 把 k 变成列向量，以便和 n 的行向量做矩阵乘
k_column = k_values[:, None]            # shape = (6,1)

# 3️⃣ 把 n 变成行向量（其实本来就是）但我们显式写出来
n_row = n[None, :]                      # shape = (1,N)

# 4️⃣ 构造旋转因子表 W[k, n] = e^{-j * 2π * k * n / N}
W = np.exp(-2j * np.pi * k_column * n_row / N)

# 需要展示的频率集合（这里只展示 k1 和 k2）
ks = [k1, k2]

# ---------------- Prepare figure layout ----------------
plt.style.use('dark_background')
fig = plt.figure(figsize=(12, 6))
gs = fig.add_gridspec(2, 2, width_ratios=[1, 1], height_ratios=[1, 1], wspace=0.3, hspace=0.35)

ax_signal = fig.add_subplot(gs[0, 0])
ax_phasor = fig.add_subplot(gs[1, 0])
ax_bar = fig.add_subplot(gs[:, 1])

# ---------- Left-top: original signal ----------
ax_signal.plot(n, x, color='#9be7ff', lw=2)
ax_signal.set_title("原始时域信号：x[n] = cos(2π·3·n/N)", fontsize=11)
ax_signal.set_xlabel("n (sample index)")
ax_signal.set_ylabel("Amplitude")
ax_signal.set_xlim(0, N-1)
ax_signal.set_ylim(-1.3, 1.3)
ax_signal.grid(alpha=0.2)

# ---------- Left-bottom: phasors ----------
ax_phasor.set_aspect('equal', 'box')
ax_phasor.set_xlim(-1.5, 1.5)
ax_phasor.set_ylim(-1.5, 1.5)
ax_phasor.set_title("旋转相量（phasors）: 当前向量（细），累积向量（粗）", fontsize=11)
ax_phasor.grid(alpha=0.2)

# unit circle
theta = np.linspace(0, 2*np.pi, 300)
ax_phasor.plot(np.cos(theta), np.sin(theta), color='#444444', lw=1, linestyle='--')

# Artists placeholders
phasor_lines = {}         # 当前短箭头 (for each k)
cumsum_lines = {}         # 累积粗箭头 (for each k)
trail_scatters = {}       # 轨迹点
colors = {k1: '#ff6b6b', k2: '#6bffb3'}  # k1 红，k2 绿

for k in ks:
    # current phasor as a line from origin to point
    line, = ax_phasor.plot([0, 0], [0, 0], lw=2, color=colors[k], alpha=0.9)
    phasor_lines[k] = line
    # cumulative vector (thicker)
    c_line, = ax_phasor.plot([0, 0], [0, 0], lw=4, color=colors[k], alpha=0.9)
    cumsum_lines[k] = c_line
    # trail scatter
    sc = ax_phasor.scatter([], [], s=20, color=colors[k], alpha=0.25)
    trail_scatters[k] = sc

# label text for magnitudes
mag_texts = {}
for idx, k in enumerate(ks):
    mag_texts[k] = ax_phasor.text(-1.4, 1.25 - 0.15*idx, f"k={k} |X|=0.00", color=colors[k], fontsize=10)

# ---------- Right: live spectrum bars ----------
bar_x = np.arange(len(ks))
bars = ax_bar.bar(bar_x, [0]*len(ks), color=[colors[k] for k in ks], alpha=0.7)
ax_bar.set_xticks(bar_x)
ax_bar.set_xticklabels([f"k={k}" for k in ks])
ax_bar.set_ylim(0, 1.05)   # 显示归一化幅度 [0,1]
ax_bar.set_title("实时频谱（归一化幅值 |X[k]| / N）", fontsize=11)
ax_bar.set_ylabel("Normalized |X[k]|")
ax_bar.grid(axis='y', alpha=0.2)

# numeric labels on bars
bar_labels = [ax_bar.text(i, 0, "0.00", ha='center', va='bottom', color='white') for i in range(len(ks))]

# ---------------- Animation state ----------------
cumulative = {k: 0+0j for k in ks}     # 累积复数和
trails = {k: [] for k in ks}           # 保存轨迹点值（复数）

# ---------------- Update function ----------------
def update(frame):
    # frame corresponds to time index n[frame]
    idx = frame % N
    for k in ks:
        # current phasor: x[n] * e^{-j*2π*k*n/N}
        ph = x[idx] * np.exp(-2j * np.pi * k * idx / N)
        # append to cumulative
        cumulative[k] += ph
        trails[k].append(ph)

        # update current phasor line (short arrow)
        ph_line = phasor_lines[k]
        ph_line.set_data([0, ph.real], [0, ph.imag])

        # update cumulative vector line (thick)
        csum = cumulative[k]
        c_line = cumsum_lines[k]
        c_line.set_data([0, csum.real], [0, csum.imag])

        # update trail scatter (show last up to 40 points for clarity)
        trail_recent = trails[k][-40:]
        if trail_recent:
            xs = [p.real for p in trail_recent]
            ys = [p.imag for p in trail_recent]
        else:
            xs, ys = [], []
        trail_scatters[k].set_offsets(np.c_[xs, ys])

        # update magnitude text (normalized by N)
        mag_norm = abs(csum) / N
        mag_texts[k].set_text(f"k={k}  |X|/N = {mag_norm:.3f}")

    # update bars (normalized magnitudes)
    for i, k in enumerate(ks):
        mag_norm = abs(cumulative[k]) / N
        bars[i].set_height(mag_norm)
        bar_labels[i].set_text(f"{mag_norm:.3f}")
        bar_labels[i].set_y(mag_norm + 0.02)
        # highlight the bar of the currently-advancing sample? instead highlight larger one
        if mag_norm == max(abs(cumulative[k2])/N for k2 in ks):
            bars[i].set_alpha(1.0)
            bars[i].set_edgecolor('white')
            bars[i].set_linewidth(1.5)
        else:
            bars[i].set_alpha(0.6)
            bars[i].set_edgecolor(None)
            bars[i].set_linewidth(0.0)

    # title with frame
    ax_phasor.set_xlabel(f"n = {idx} / {N-1}")
    return list(phasor_lines.values()) + list(cumsum_lines.values()) + list(trail_scatters.values()) + list(bars) + list(bar_labels) + list(mag_texts.values())

# ---------------- Run animation ----------------
ani = FuncAnimation(fig, update, frames=N, interval=interval_ms, blit=False, repeat=False)

plt.show()
